import cv2
import random
from ultralytics import YOLO
from ultralytics import solutions


class Detector():
    def init_model(self,model_path):
        self.model = YOLO(model_path)
        self.names = self.model.names
        self.colors = {name: (random.randint(0, 255), random.randint(0, 255), random.randint(0, 255)) for name in
                  self.names.values()}
        self.heatmap = solutions.Heatmap(
            show=False,
            model = model_path,
            colormap=cv2.COLORMAP_PARULA,
            classes=[0, 2]
        )

    def detect(self,image,conf=0.3,iou=0.3):
        detections = []
        results = self.model.predict(image,conf=conf,iou=iou)[0]
        boxes = results.boxes.data.tolist()
        for obj in boxes:
            left, top, right, bottom = int(obj[0]), int(obj[1]), int(obj[2]), int(obj[3])
            confidence = obj[4]
            label = self.model.names[int(obj[5])]
            detections.append((left, top, right, bottom,label,round(confidence,2)))
        return detections

    def drawboxs(self, image, detections):
        # 为每个标签名生成一个对应的颜色
        font = cv2.FONT_HERSHEY_SIMPLEX
        box_thickness = 2
        font_thickness = 2

        max_font_scale = 0.7
        min_font_scale = 1.5
        tl = round(0.002 * (image.shape[0] + image.shape[1]) / 2) + 1  # line/font thickness

        for x1, y1,x2,y2,label, confidence in detections:
            start_point = (x1, y1)
            end_point = (x2, y2)
            color = self.colors.get(label, (0, 255, 0))
            cv2.rectangle(image, start_point, end_point, color, box_thickness)

            # 根据框的高度动态调整字体大小
            box_height = y2 - y1
            font_scale = box_height / 100  # 字体大小与框高度成比例
            font_scale = max(min_font_scale, min(max_font_scale, font_scale))

            text = f'{label} {confidence:.2f}'
            (text_width, text_height), baseline = cv2.getTextSize(text, font, font_scale, font_thickness)
            text_origin = (start_point[0], start_point[1])
            background_tl = (text_origin[0], text_origin[1] - text_height - baseline)
            background_br = (text_origin[0] + text_width, text_origin[1])
            cv2.rectangle(image, background_tl, background_br, color, cv2.FILLED)
            cv2.putText(image, text, (text_origin[0], text_origin[1] - baseline), font, font_scale, (255, 255, 255), font_thickness)
        return image


# if __name__ == '__main__':
#     detect = Detector()
#     detect.init_model('model_weight/hatV5_best.pt')
#     img = cv2.imread("test_res/test_img/hard_hat_workers65.png")
#     res = detect.detect(img,0.25,0.45)
#     img = detect.drawboxs(img,res)
#     cv2.imshow("fdsf",img)
#     cv2.waitKey(0)

# if __name__ == '__main__':
#     from ultralytics import YOLO
#
#     # Load a model
#     model = YOLO("yolo11n.yaml")  # build a new model from YAML
#     model = YOLO("yolo11n.pt")  # load a pretrained model (recommended for training)
#     model = YOLO("yolo11n.yaml").load("yolo11n.pt")  # build from YAML and transfer weights
#
#     # Train the model
#     results = model.train(data="/root/autodl-tmp/lxh/Ultralytics/data.yaml", epochs=100, imgsz=640)
